package cn.bugstack.infrastructure.comm;

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * 加权随机选择器，用于按概率抽取元素
 * @param <T> 元素类型，例如：Integer（奖品ID）、String（策略名称）等
 */
public class WeightedRandomSelector<T> {

    private final List<BigDecimal> cumulativeRates = new ArrayList<>();
    private final List<T> items = new ArrayList<>();
    private BigDecimal totalRate = BigDecimal.ZERO;

    /**
     * 添加一个带权重的元素
     * @param rate 权重（概率），如 0.1 表示 10%
     * @param item 对应的元素，如奖品 ID
     */
    public void add(BigDecimal rate, T item) {
        if (rate == null || rate.compareTo(BigDecimal.ZERO) < 0) {
            throw new IllegalArgumentException("权重必须是非负数");
        }
        if (item == null) {
            throw new IllegalArgumentException("元素不能为空");
        }

        totalRate = totalRate.add(rate);
        cumulativeRates.add(totalRate);
        items.add(item);
    }

    /**
     * 按照权重随机选取一个元素
     * @return 被选中的元素
     */
    public T next() {
        if (totalRate.compareTo(BigDecimal.ZERO) <= 0) {
            throw new IllegalStateException("权重总和不能为0");
        }

        BigDecimal randomRate = new BigDecimal(Math.random()).multiply(totalRate);
        int index = Collections.binarySearch(cumulativeRates, randomRate, BigDecimal::compareTo);
        if (index < 0) {
            index = -index - 1;
        }

        return items.get(index);
    }

    /**
     * 获取当前所有项的数量
     */
    public int size() {
        return items.size();
    }

    /**
     * 清空所有已添加的项
     */
    public void clear() {
        cumulativeRates.clear();
        items.clear();
        totalRate = BigDecimal.ZERO;
    }
}
